import pandas as pd
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score
import numpy as np

from fib_config import close_fib_file, y_column, y_scale, tezheng, data_split_rate
from buy_flow import buy_flow_v2

"""
特征值放大10倍取整
3日、5日收益
准确率 0.8
收益率 20%左右
todo 打印出时间和股票来具体分析
"""


def decision():
    df = pd.read_csv(close_fib_file)
    print(df.columns)
    print(df.shape)
    ### 删除NAN，看下数据的正确性
    df = df.dropna(axis=0, how='any')
    x = df[tezheng].values
    y = np.around((df[y_column] / y_scale).values, 0)
    data_split = int(df.index.size * data_split_rate)
    # x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.5)
    x_train, x_test = x[: data_split], x[data_split:]
    y_train, y_test = y[: data_split], y[data_split:]
    print("X_train_shape:", x_train.shape, " y_train_shape:", y_train.shape)
    print("X_test_shape:", x_test.shape, "  y_test_shape:", y_test.shape)

    from collections import Counter

    print('Counter(data)\n', Counter(np.around(y_test)))
    # max_depth 5 或 6
    # 剪纸
    # 随机森林
    # dec = DecisionTreeClassifier(max_depth=6)
    # dec.fit(x_train, y_train.astype('int'))
    dec = XGBClassifier(objective="reg:linear",
                        eval_metric="rmse",
                        max_depth=6,
                        learning_rate=0.15)
    dec.fit(x_train, np.around(y_train.astype('int')))
    # dec.save_model("stcok_xgb_1.1.model")
    # dec.load_model("stcok_xgb_1.model")
    # print("----", dec.score(x_train, y_train.astype('int')))
    print("----", dec.score(x_test, np.around(y_test).astype('int')))

    y_pred = dec.predict(x_test)
    accuracy = accuracy_score(np.around(y_test.astype('int')), y_pred)
    print("accuarcy: %.2f%%" % (accuracy * 100.0))
    # buy_flow(x_test, y_test, dec)
    buy_flow_v2(df[data_split:], dec)
    print("accuarcy: %.2f%%" % (accuracy * 100.0))
    print("------------------end")

    import matplotlib.pyplot as plt

    from xgboost import plot_importance

    fig, ax = plt.subplots(figsize=(10, 15))
    plot_importance(dec, height=0.5, max_num_features=64, ax=ax)
    plt.show()
    # 参数重要性
    important = dict(zip(tezheng, dec.feature_importances_))
    d_order = sorted(important.items(), key=lambda x: x[1], reverse=False)
    for i in d_order:
        print(i)
    print("end")

    # export_graphviz(dec, feature_names=tezheng, out_file="./tree.dot")
    # dot -Tpdf tree.dot -o tree.pdf


def load_decision():
    dec = XGBClassifier()
    dec.load_model("stcok_xgb.model")


if __name__ == '__main__':
    decision()
